//
// Created by 11478 on 2024/7/4.
//

#include "UDP.h"

UDP::UDP(int port, std::function<void(const char *, int, int)> recvFunction)
        : port(port), recvFunction(recvFunction), stop(false) {
    sockfd = socket(AF_INET, SOCK_DGRAM, 0);

    memset(&servaddr, 0, sizeof(servaddr));

    servaddr.sin_family = AF_INET;
    servaddr.sin_addr.s_addr = INADDR_ANY;
    servaddr.sin_port = htons(port);

    if (bind(sockfd, (const struct sockaddr *) &servaddr, sizeof(servaddr)) < 0) {
        perror("bind failed");
        close(sockfd);
        exit(EXIT_FAILURE);
    }

    std::cout << "UDP server initialized on port " << port << std::endl;
}

UDP::~UDP() {
    stop = true;
    if (receiverThread.joinable()) {
        receiverThread.join();
    }
    close(sockfd);
}

void UDP::start() {
    receiverThread = std::thread(&UDP::receive, this);
}

void UDP::receive() {
    char buffer[1024];
    struct sockaddr_in cliaddr;
    memset(&cliaddr, 0, sizeof(cliaddr));
    socklen_t len = sizeof(cliaddr);

    while (!stop) {
        int n = recvfrom(sockfd, (char *) buffer, 1024, MSG_WAITALL, (struct sockaddr *) &cliaddr, &len);
        if (n > 0) {
            buffer[n] = '\0';
            recvFunction(buffer, n, port);
        }
    }
}

UDPSender::UDPSender(const std::string &address, int port) {
    sockfd = socket(AF_INET, SOCK_DGRAM, 0);
    if (sockfd < 0) {
        perror("socket creation failed");
        exit(EXIT_FAILURE);
    }

    memset(&destAddr, 0, sizeof(destAddr));
    destAddr.sin_family = AF_INET;
    destAddr.sin_port = htons(port);
    inet_pton(AF_INET, address.c_str(), &destAddr.sin_addr);
}

UDPSender::~UDPSender() {
    close(sockfd);
}

void UDPSender::send(const std::vector<char>& message) {
    sendto(sockfd, message.data(), message.size(), 0, (const struct sockaddr*)&destAddr, sizeof(destAddr));
}

void UDPManager::addUDP(int port, std::function<void(const char *, int, int)> recvFunction) {
    udps.push_back(std::unique_ptr<UDP>(new UDP(port, recvFunction)));
}

void UDPManager::addSender(const std::string &address, int port) {
    senders.push_back(std::unique_ptr<UDPSender>(new UDPSender(address, port)));
}

void UDPManager::startAll() {
    for (auto &udp: udps) {
        udp->start();
    }
}

void UDPManager::sendToAll(std::vector<char>& message) {
    for (auto &sender: senders) {
        sender->send(message);
    }
}
